import argparse
import os
import random
import sys
import joblib
import torch

import numpy as np
import pandas as pd
from torch.utils.data import DataLoader
from orbitP.script.DataLoader import orbitPSULTDataset
from orbitP.script.util import clean_directory, init_args, splitDataset
from orbitP.script.train import train_SULT_Seq2Seq
from orbitP.script import config

np.random.seed(42)
random.seed(42)
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)

def main(args):
    if config.resumeModel == 0 and config.usePTModel == 0:
        clean_directory()
    df_obserData = pd.read_csv(config.dataSetDir + 'df_obsData.csv')
    df_prdData = pd.read_csv(config.dataSetDir + 'df_prdData.csv')
    df_stampObs = pd.read_csv(config.dataSetDir + 'df_stampObs.csv')
    df_stampPrd = pd.read_csv(config.dataSetDir + 'df_stampPrd.csv')
    obsData = df_obserData.to_numpy(dtype=np.float32)
    prdData = df_prdData.to_numpy(dtype=np.float32)
    stampObs = df_stampObs.to_numpy(dtype=np.float32)
    stampPrd = df_stampPrd.to_numpy(dtype=np.float32)

    obsData, prdData, stampObs, stampPrd = splitDataset(obsData, prdData, stampObs, stampPrd, runType='train', sca=True)
    obsData_train = obsData[0];obsData_val = obsData[1];obsData_test = obsData[2]
    prdData_train = prdData[0];prdData_val = prdData[1];prdData_test = prdData[2]
    stampObs_train = stampObs[0];stampObs_val = stampObs[1];stampObs_test = stampObs[2]
    stampPrd_train = stampPrd[0];stampPrd_val = stampPrd[1];stampPrd_test = stampPrd[2]

    train_dataset = orbitPSULTDataset(obsData=obsData_train,prdData=prdData_train,stampObs=stampObs_train,stampPrd=stampPrd_train,training_length=config.training_length, predicting_length=config.predicting_length, forecast_window=config.forecast_window, days=obsData_train.shape[0]/config.training_length)
    val_dataset = orbitPSULTDataset(obsData=obsData_val,prdData=prdData_val,stampObs=stampObs_val,stampPrd=stampPrd_val,training_length=config.training_length, predicting_length=config.predicting_length, forecast_window=config.forecast_window, days=prdData_val.shape[0]/config.predicting_length)
    test_dataset = orbitPSULTDataset(obsData=obsData_test,prdData=prdData_test,stampObs=stampObs_test,stampPrd=stampPrd_test,training_length=config.training_length, predicting_length=config.predicting_length, forecast_window=config.forecast_window, days=prdData_test.shape[0]/config.predicting_length)

    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=0, shuffle=True, pin_memory=True)
    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, num_workers=0, shuffle=False, pin_memory=True)
    test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=0, shuffle=False, pin_memory=True)
    model = train_SULT_Seq2Seq(train_dataloader, val_dataloader, args)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="itransformer")
    args = parser.parse_args()
    args = init_args(args)

    main(args)

